# -*- coding:utf-8 -*-

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import argparse
import re
parser = argparse.ArgumentParser(description="Data Augmenation")
parser.add_argument("--DA", type=bool, default=False)
args = parser.parse_args()
pattern = '\d+'
#sets=[('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
sets=[('2020', 'train'),('2020', 'val')]

#classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
# classes=['U盘,7','书本,5','充电器,3','充电宝,3','其他垃圾,2','勺子,2','包,4','医用面罩,3','卫生用品,2','卫生纸,2','厨余垃圾,8','发夹,2','口香糖,8','叶子,2','吸管,7','啤酒瓶盖,2','坚果壳,8','塑料,7','塑料刀叉,7','塑料包装,7','塑料化妆品,7','塑料快餐盒,7','塑料手套,7','塑料晾衣架,7','塑料杯,7','塑料玩具,7','塑料瓶盖,7','塑料盘,7','塑料碗,7','塑料管,7','塑料篮子,7','塑料粘性钩,7','塑料花盆,7','塑料袋,7','塑料饮料瓶,1','奶茶杯,7','布手套,4','布面罩,4','带塑料钥匙,7','帽子,4','快递纸带,5','手机,6','手机壳,2','手表,6','打火机,2','报纸,5','插头,7','旧枕头,4','旧衣服,4','易拉罐,1','有害垃圾,3','木筷子,2','果核,8','果皮,8','果肉,8','棉签,3','毛巾,2','毛绒玩具,4','泡沫塑料,7','洗发水瓶,7','烟灰缸,2','烟盒,5','烟蒂,2','牙刷,7','牙签,2','牙膏,2','牛奶盒,5','玻璃,2','玻璃杯,2','玻璃瓶,2','珠宝,2','电池,3','电线,7','电饭煲,6','盘子,2','砧板,2','碗,2','笔,7','纸,5','纸快餐盒,5','纸提袋,5','纸杯,5','纸板,5','纸盒,5','织物,4','绳子,2','羽毛,2','耳机,7','茶叶渣,8','药物,3','菜叶,8','菜根,8','蔬菜,8','蛋壳,8','袜子,4','调料瓶,2','身份证,7','软膏,3','遥控器,7','金属,6','金属刀叉,6','金属化妆品,6','金属晾衣架,6','金属筷子,6','金属篮子,6','金属粘性钩,6','金属钥匙,6','金属食品罐,6','金属饰品,6','钱包,2','铁锅,6','银行卡,7','锅铲,6','锡纸,2','键盘,7','陶瓷花盆,2','陶锅,2','面包,8','鞋子,4','鞋带,4','食用油桶,7','骨头,8','鼠标,7']
classes = ['holothurian', 'echinus','scallop','starfish']

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(year, image_id):
    in_file = open('/home/chenzl/my/objection_detection_under_water/VOC%s/Annotation/%s.xml'%(year, image_id),encoding="utf-8")
    out_file = open('/home/chenzl/my/objection_detection_under_water/VOC%s/labels/%s.txt'%(year, image_id), 'w')

    #============================Data Augmentation===========================
    #========================================================================
    if args.DA ==True:
        out_file_LR90 = open('/usr/VOCdevkit/VOC%s/labels/%sLR90.txt' % (year, image_id), 'w')
        out_file_RR90 = open('/usr/VOCdevkit/VOC%s/labels/%sRR90.txt' % (year, image_id), 'w')
        out_file_CTB = open('/usr/VOCdevkit/VOC%s/labels/%sCTB.txt' % (year, image_id), 'w')
        out_file_CLR = open('/usr/VOCdevkit/VOC%s/labels/%sCLR.txt' % (year, image_id), 'w')
    #========================================================================

    try:
        tree=ET.parse(in_file)
    except Exception:
        print(image_id)

    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        #===========没用labelimg标注======
        # difficult = obj.find('difficult').text
        # cls = obj.find('name').text
        # if cls not in classes or int(difficult)==1:
        #     continue
        #=================================
        # difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
        #================================Data Augmentation==========================
        '''
                                          DA RULE
                                      ORI (x, y, w, h)
                                      CTB (x, 1-y, w, h)
                                      CLR (1-x, y, w, h)
                                      RR90(1-y, x, h, w)
                                      LR90(y, 1-x, h, w)
          '''
        #====================bb[0]:x, bb[1]:y, bb[2]:w, bb[3]:h=====================
        if args.DA == True:
            LR90 = [bb[1], 1-bb[0], bb[3], bb[2]]
            out_file_LR90.write(str(cls_id) + " " + " ".join([str(a) for a in LR90]) + '\n')

            RR90 = [1-bb[1], bb[0], bb[3], bb[2]]
            out_file_RR90.write(str(cls_id) + " " + " ".join([str(a) for a in RR90]) + '\n')

            CTB = [bb[0], 1-bb[1], bb[2], bb[3]]
            out_file_CTB.write(str(cls_id) + " " + " ".join([str(a) for a in CTB]) + '\n')

            CLR = [1-bb[0], bb[1], bb[2], bb[3]]
            out_file_CLR.write(str(cls_id) + " " + " ".join([str(a) for a in CLR]) + '\n')
        #===========================================================================

#wd = getcwd()
os.system('rm /home/chenzl/my/objection_detection_under_water/VOC2020/2020_train.txt')
os.system('rm /home/chenzl/my/objection_detection_under_water/VOC2020/2020_val.txt')
os.system('rm -R /home/chenzl/my/objection_detection_under_water/VOC2020/labels/*.txt')
for year, image_set in sets:
    if not os.path.exists('/home/chenzl/my/objection_detection_under_water/VOC%s/labels/'%(year)):
        os.makedirs('/home/chenzl/my/objection_detection_under_water/VOC%s/labels/'%(year))
    image_ids = open('/home/chenzl/my/objection_detection_under_water/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('/home/chenzl/my/objection_detection_under_water/%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        try:
            list_file.write('/home/chenzl/my/objection_detection_under_water/VOC%s/JPEGImages/%s.jpg\n' % (year, image_id))
            # if len(re.findall('\d+', image_id)[0]) == len(image_id):
                # list_file.write('/home/user/my/rubbishDataset/JPEGImages/%s.jpg\n' %(image_id)) #remeber to modifier the cfg
            convert_annotation(year, image_id)
        except Exception:
            print(image_id)
    list_file.close()
    os.system(r'chown -R user /home/chenzl/my/objection_detection_under_water')

#os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
#os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")

